package com.example.ai.service;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationOutput;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.example.ai.entity.AIRequest;
import com.example.ai.entity.ProviderEnum;
import io.reactivex.Flowable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.Arrays;
import java.util.Collections;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

@Service
public class AIService {

    /**
     * 日志
     */
    private static final Logger logger = LoggerFactory.getLogger(AIService.class);

    @Value("${ai.api-key}")
    private String apiKey;

    @Value("${ai.temperature}")
    private float temperature;

    /**
     * 历史对话
     */
    @Autowired
    private ChatHistoryService chatHistoryService;


    /**
     * 大模型文本
     *
     * @param request  请求参数
     * @param provider 大模型类型
     * @return 结果
     */
    public String generateText(AIRequest request, String provider) {

        ProviderEnum providerEnum = ProviderEnum.getInstance(provider);

        Generation gen = new Generation();
        Message userMsg = Message.builder().role(Role.USER.getValue()).content(request.getPrompt()).build();
        GenerationParam param = GenerationParam.builder().apiKey(apiKey).model(request.getModel() != null ? request.getModel() : providerEnum.getDefModel()).messages(Arrays.asList(userMsg))
                // 不可以设置为"text"
                .resultFormat(GenerationParam.ResultFormat.MESSAGE).temperature(ObjectUtil.isNotEmpty(request.getTemperature()) ? request.getTemperature() : temperature).build();

        String content = StrUtil.EMPTY;
        try {
            GenerationResult result = gen.call(param);
            if (ProviderEnum.DEEPSEEK.getValue().equals(providerEnum.getValue())) {
                content = "思考过程：\n";
                content = content + result.getOutput().getChoices().get(0).getMessage().getReasoningContent() + "\n";
            }
            content = content + "回复内容：";
            content = content + result.getOutput().getChoices().get(0).getMessage().getContent();
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            // 使用日志框架记录异常信息
            logger.error("An error occurred while calling the generation service: {}", e.getMessage());
        }
        return content;
    }

    private static boolean isFirstPrint = true;

    /**
     * 流式返回
     *
     * @param request  请求参数
     * @param provider 大模型类型
     * @return 结果
     */
    public Flux<ServerSentEvent<String>> streamGenerate1(AIRequest request, String provider) {
        ProviderEnum providerEnum = ProviderEnum.getInstance(provider);
        Generation gen = new Generation();
        Message userMsg = Message.builder().role(Role.USER.getValue()).content(request.getPrompt()).build();

        GenerationParam param = GenerationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(apiKey)
                .model(request.getModel() != null ? request.getModel() : providerEnum.getDefModel())
                .messages(Arrays.asList(userMsg))
                // 不可以设置为"text"
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();
        Flowable<GenerationResult> generationResultFlowable = null;
        try {
            generationResultFlowable = gen.streamCall(param);
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            logger.error("流式返回报错，报错信息: {}", e.getMessage());
        }
        StringBuilder reasoningContent = new StringBuilder();
        StringBuilder finalContent = new StringBuilder();
        return Flux.from(generationResultFlowable).map(message -> {
            //思考过程
            String reasoning = message.getOutput().getChoices().get(0).getMessage().getReasoningContent();
            String contentMessage = StrUtil.EMPTY;
            if (StrUtil.isNotBlank(reasoning)) {
                reasoningContent.append(reasoning);
                if (isFirstPrint) {
                    isFirstPrint = false;
                    contentMessage = "\n====================思考过程====================";
                } else if (StrUtil.isNotBlank(reasoningContent.toString())) {
                    contentMessage = reasoningContent.toString();
                }
            }
            String content = message.getOutput().getChoices().get(0).getMessage().getContent();
            if (StrUtil.isNotBlank(content)) {
                finalContent.append(content);
                if (!isFirstPrint) {
                    isFirstPrint = true;
                    contentMessage = "\n====================完整回复====================";
                } else if (StrUtil.isNotBlank(finalContent.toString())) {
                    contentMessage = finalContent.toString();
                }
            }
            return ServerSentEvent.builder(contentMessage)
                    //.event("ai-message")// 事件类型标识
                    .build();
        });
    }

    public Flux<ServerSentEvent<String>> streamGenerate(AIRequest request, String provider, String userId) {
        return Flux.defer(() -> {
            try {
                ProviderEnum providerEnum = ProviderEnum.getInstance(provider);
                Generation gen = new Generation();
                Message userMsg = Message.builder()
                        .role(Role.USER.getValue())
                        .content(request.getPrompt())
                        .build();

                GenerationParam param = GenerationParam.builder()
                        .apiKey(apiKey)
                        .model(request.getModel() != null ? request.getModel() : providerEnum.getDefModel())
                        .messages(Collections.singletonList(userMsg))
                        .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                        .incrementalOutput(true)
                        .build();

                Flowable<GenerationResult> generationResultFlowable = gen.streamCall(param);

                // 使用原子变量处理状态（线程安全）
                AtomicBoolean isFirstPrint = new AtomicBoolean(true);
                StringBuilder reasoningContent = new StringBuilder();
                StringBuilder finalContent = new StringBuilder();
                AtomicReference<Integer> totalTokens = new AtomicReference<>(0);
                return Flux.from(generationResultFlowable)
                        .map(message -> {
                            // 安全获取消息内容
                            GenerationOutput output = Optional.ofNullable(message)
                                    .map(GenerationResult::getOutput)
                                    .orElseThrow(() -> new IllegalStateException("Empty message output"));

                            GenerationOutput.Choice choice = Optional.ofNullable(output.getChoices())
                                    .filter(choices -> !choices.isEmpty())
                                    .map(choices -> choices.get(0))
                                    .orElseThrow(() -> new IllegalStateException("No choices available"));

                            Message msg = Optional.ofNullable(choice.getMessage())
                                    .orElseThrow(() -> new IllegalStateException("Empty message"));

                            // 处理思考过程
                            String reasoning = msg.getReasoningContent();
                            String content = msg.getContent();

                            String contentMessage = processContent(
                                    reasoning,
                                    content,
                                    reasoningContent,
                                    finalContent
                            );
                            String event = "ai-message";
                            if (StringUtils.hasText(reasoning)) {
                                event = "reasoning-message";
                            }
                            totalTokens.set(message.getUsage().getTotalTokens());

                            return ServerSentEvent.builder(contentMessage)
                                    .event(event)
                                    .build();
                        })// 关键修复点：过滤空事件
                        .filter(event ->
                                !event.data().isBlank() &&
                                        !event.data().equals("[DONE]")
                        )
                        // 添加明确的流结束标识
                        .concatWith(Flux.just(
                                ServerSentEvent.<String>builder()
                                        .event("end")
                                        .data("[DONE]")
                                        .build()
                        ))
                        .onErrorResume(e -> {
                            logger.error("流式处理异常", e);
                            return Flux.error(new Throwable("AI服务处理失败"));
                        }).doOnComplete(()->{
                            chatHistoryService.saveMessagePair(userId, request.getSessionId(), request.getPrompt(), finalContent.toString(), totalTokens.get());
                        });
            } catch (ApiException | NoApiKeyException | InputRequiredException e) {
                logger.error("初始化流失败", e);
                return Flux.error(new Throwable("AI服务初始化失败"));
            }
        });
    }

    private String processContent(String reasoning,
                                  String content,
                                  StringBuilder reasoningContent,
                                  StringBuilder finalContent) {
        String contentMessage = StrUtil.EMPTY;

        // 处理思考内容
        if (StringUtils.hasText(reasoning)) {
            reasoningContent.append(reasoning);
            contentMessage = reasoning.toString();
        }

        // 处理最终内容
        if (StringUtils.hasText(content)) {
            finalContent.append(content);
            contentMessage = content.toString();
        }
        return StringUtils.hasText(contentMessage) ? contentMessage : "";
    }
}
